import torch, os, cv2
from model.model import parsingNet
import matplotlib.pyplot as plt
from utils1.common import merge_config
from utils1.dist_utils import dist_print
import torch
import scipy.special, tqdm
import numpy as np
import torchvision.transforms as transforms
from data.dataset import LaneTestDataset
from data.constant import culane_row_anchor, tusimple_row_anchor
import sys
import time
from PIL import Image

project_path = 'E:/pythonprojectnew/graduation/Ultra-Fast-Lane-Detection-master'
sys.path.append(project_path)
video_path=r'E:\pythonprojectnew\graduation\CarLaneDetection-master\resources\v2.mp4'
# img_path=r'E:\pythonprojectnew\graduation\Ultra-Fast-Lane-Detection-master\CUlane\driver_100_30frame\05251548_04392.MP4\00000.JPG'
img_path=r'C:\Users\Administrator\Pictures\school.jpg'

def linear_regression(x, y):
    N = len(x)
    sumx = sum(x)
    sumy = sum(y)
    sumx2 = sum(x ** 2)
    sumxy = sum(x * y)

    A = np.mat([[N, sumx], [sumx, sumx2]])
    b = np.array([sumy, sumxy])

    return np.linalg.solve(A, b)



if __name__ == "__main__":

    cv2.namedWindow("rec_img", 0)
    cv2.resizeWindow("rec_img",1640, 590)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.backends.cudnn.benchmark = True

    args, cfg = merge_config()

    dist_print('start testing...')
    # assert cfg.backbone in ['18', '34', '50', '101', '152', '50next', '101next', '50wide', '101wide']

    cls_num_per_lane = 18
    img_w, img_h = 1640, 590
    row_anchor = culane_row_anchor
    # net = parsingNet(pretrained = False, backbone=cfg.backbone,cls_dim = (cfg.griding_num+1,cls_num_per_lane,4),
    #                 use_aux=False).cuda() # we dont need auxiliary segmentation in testing
    net = parsingNet(pretrained=False, backbone=cfg.backbone, cls_dim=(cfg.griding_num + 1, cls_num_per_lane, 4),
                     use_aux=False).to(device)  # we dont need auxiliary segmentation in testing

    state_dict = torch.load(cfg.test_model, map_location='cpu')['model']
    compatible_state_dict = {}
    for k, v in state_dict.items():
        if 'module.' in k:
            compatible_state_dict[k[7:]] = v
        else:
            compatible_state_dict[k] = v

    net.load_state_dict(compatible_state_dict, strict=False)
    net.eval()

    img_transforms = transforms.Compose([
        transforms.Resize((288, 800)),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])

    org_image = Image.open(img_path)#读取
    print("org_image",type(org_image),org_image)
    org_image = img_transforms(org_image)#转换，同时变成torch.tensor格式
    img_in = org_image.unsqueeze(0)#增加一维度
    print("img_in", type(img_in), img_in.shape)
    #定义视频格式和名字灯
    #fourcc = cv2.VideoWriter_fourcc(*'MJPG')
    #vout = cv2.VideoWriter('test' + 'avi', fourcc, 30.0, (img_w, img_h))

    #输入图片
    imgs = img_in.to(device)
    with torch.no_grad():
        out = net(imgs)
    # col_sample = np.linspace(0, 800 - 1, cfg.griding_num)
    # col_sample_w = col_sample[1] - col_sample[0]
    col_sample_w=4.015075376884422

    out_j = out[0].data.cpu().numpy()
    out_j = out_j[:, ::-1, :]
    prob = scipy.special.softmax(out_j[:-1, :, :], axis=0)
    idx = np.arange(cfg.griding_num) + 1
    idx = idx.reshape(-1, 1, 1)
    loc = np.sum(prob * idx, axis=0)
    out_j = np.argmax(out_j, axis=0)

    loc[out_j == cfg.griding_num] = 0
    out_j = loc

    # import pdb; pdb.set_trace()
    org_img = cv2.imread(img_path)
    img_rec = cv2.resize(org_img, (img_w, img_h), cv2.INTER_AREA)
    # print(np.sum(out_j[:, 1]))
    x=[]
    y=[]
    count2=0
    print("out_j",out_j)
    '''
    for i in range(out_j.shape[1]):
        if np.sum(out_j[:, i] != 0) > 2:
            for k in range(out_j.shape[0]):
                if out_j[k, i] > 0:
                    ppp = (int(out_j[k, i] * col_sample_w * img_w / 800) - 1,
                            int(img_h * (row_anchor[cls_num_per_lane - 1 - k] / 288)) - 1)
                    # x.append(int(out_j[k, i] * col_sample_w * img_w / 800) - 1)
                    # y.append(int(img_h * (row_anchor[cls_num_per_lane - 1 - k] / 288)) - 1)
                    cv2.circle(img_rec, ppp, 5, (0, 255, 0), -1)
                    print(count2)
                    # print("ppp",ppp)
    '''
    if np.sum(out_j[:, 1] != 0) > 2:
        for k in range(out_j.shape[0]):
            if out_j[k, 1] > 0:
                x.append(int(out_j[k, 1] * col_sample_w * img_w / 800) - 1)
                y.append(int(img_h * (row_anchor[cls_num_per_lane - 1 - k] / 288)) - 1)
    x = np.array(x)
    y = np.array(y)
    a0, a1 = linear_regression(y,x)
    # 生成拟合直线的绘制点
    _Y = [590]
    _X = [a0 + a1 * x for x in _Y]
    ppp = (int(_X[0]),
           int(img_h * (row_anchor[cls_num_per_lane - 1 - 0] / 288)) - 1)
    cv2.circle(img_rec, ppp, 5, (0, 255, 0), -1)
    # print(_Y,_X)
    # print(_X,_Y)
    # plt.plot( _X, _Y, 'b', linewidth=2)
    # print("xy",x,y)
    # plt.scatter(x, y)
    # plt.gca().invert_yaxis()
    # plt.show()
    cv2.imshow('rec_img', img_rec)
    cv2.waitKey(0)

    cv2.destroyAllWindows()
    #vout.write(img_rec)
    #vout.release()



